# spark_chat.py - 修复版星火对话模块
import json
import ssl
import hmac
import base64
import hashlib
import threading
import time
from urllib.parse import urlencode, urlparse
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
import logging

logger = logging.getLogger(__name__)

# 检查WebSocket可用性
try:
    import websocket

    WEBSOCKET_AVAILABLE = True
    logger.info("WebSocket模块可用")
except ImportError:
    WEBSOCKET_AVAILABLE = False
    logger.warning("WebSocket模块不可用")

from config import SPARK_APPID, SPARK_APIKey, SPARK_APISecret, SPARK_URL, SPARK_DOMAIN


class SparkWs:
    """星火大模型WebSocket客户端"""

    def __init__(self, appid, apikey, apisecret, spark_url, domain, system_prompt):
        self.appid = appid
        self.apikey = apikey
        self.apisecret = apisecret
        self.spark_url = spark_url
        self.domain = domain
        self.system_prompt = system_prompt
        self.history = [{"role": "system", "content": system_prompt}]
        self.answer = ""
        self.ws_url = self._build_url()
        self._connection_error = False

    def _build_url(self):
        """构建WebSocket连接URL"""
        try:
            u = urlparse(self.spark_url)
            host, path = u.hostname, u.path
            date = format_date_time(mktime(datetime.now().timetuple()))

            # 构建签名
            sign_origin = f"host: {host}\ndate: {date}\nGET {path} HTTP/1.1"
            sign_sha = hmac.new(
                self.apisecret.encode(),
                sign_origin.encode(),
                digestmod=hashlib.sha256
            ).digest()
            signature = base64.b64encode(sign_sha).decode()

            # 构建授权字符串
            auth_str = f'api_key="{self.apikey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature}"'
            auth_base64 = base64.b64encode(auth_str.encode()).decode()

            return f"{self.spark_url}?{urlencode({'authorization': auth_base64, 'date': date, 'host': host})}"
        except Exception as e:
            logger.error(f"构建WebSocket URL失败: {e}")
            return None

    def send_prompt(self, prompt: str) -> str:
        """发送提示并获取回复"""
        if not WEBSOCKET_AVAILABLE:
            logger.warning("WebSocket不可用，使用备用回复")
            return self._get_fallback_response(prompt)

        if not self.ws_url:
            logger.error("WebSocket URL无效，使用备用回复")
            return self._get_fallback_response(prompt)

        self.answer = ""
        self._connection_error = False

        # 使用同步方式处理WebSocket连接
        result = self._sync_websocket_request(prompt)

        if not result or self._connection_error:
            logger.warning("星火大模型响应异常，使用备用回复")
            return self._get_fallback_response(prompt)

        return result

    def _sync_websocket_request(self, prompt):
        """同步WebSocket请求"""
        try:
            # 使用事件来同步等待结果
            import threading
            result_event = threading.Event()
            result_container = {"answer": ""}

            def on_open(ws):
                try:
                    # 构建消息
                    if not prompt:
                        messages = [
                            {"role": "system", "content": "你是一位专业面试官，请用中文进行面试。"},
                            {"role": "user", "content": "请开始面试，问一个开场问题。"}
                        ]
                    else:
                        # 只保留最近的对话，避免消息过长
                        recent_history = self.history[-2:] if len(self.history) > 2 else self.history
                        messages = recent_history + [{"role": "user", "content": prompt}]

                    # 构建请求体
                    body = {
                        "header": {"app_id": self.appid, "uid": "user"},
                        "parameter": {"chat": {"domain": self.domain, "temperature": 0.7, "max_tokens": 2048}},
                        "payload": {"message": {"text": messages}}
                    }

                    logger.info("发送消息到星火大模型...")
                    ws.send(json.dumps(body))

                except Exception as e:
                    logger.error(f"发送消息失败: {e}")
                    self._connection_error = True
                    result_event.set()

            def on_message(ws, message):
                try:
                    data = json.loads(message)
                    if data["header"]["code"] != 0:
                        logger.error(f"星火响应错误: {data['header']['message']}")
                        self._connection_error = True
                        result_event.set()
                        return

                    content = data["payload"]["choices"]["text"][0].get("content", "")
                    result_container["answer"] += content

                    if data["payload"]["choices"]["status"] == 2:
                        logger.info("对话完成")
                        result_event.set()

                except Exception as e:
                    logger.error(f"处理消息失败: {e}")
                    self._connection_error = True
                    result_event.set()

            def on_error(ws, error):
                logger.error(f"WebSocket错误: {error}")
                self._connection_error = True
                result_event.set()

            def on_close(ws, close_status_code, close_msg):
                logger.info("WebSocket连接关闭")
                result_event.set()

            # 创建WebSocket连接
            ws = websocket.WebSocketApp(
                self.ws_url,
                on_open=on_open,
                on_message=on_message,
                on_error=on_error,
                on_close=on_close
            )

            # 在单独的线程中运行WebSocket
            def run_websocket():
                ws.run_forever(
                    sslopt={"cert_reqs": ssl.CERT_NONE},
                    ping_interval=30,
                    ping_timeout=10
                )

            ws_thread = threading.Thread(target=run_websocket, daemon=True)
            ws_thread.start()

            # 等待结果，最多30秒
            if result_event.wait(timeout=30):
                if not self._connection_error and result_container["answer"]:
                    # 更新历史记录
                    if prompt:
                        self.history.append({"role": "user", "content": prompt})
                        self.history.append({"role": "assistant", "content": result_container["answer"]})

                    # 限制历史记录长度
                    if len(self.history) > 10:
                        self.history = self.history[:1] + self.history[-8:]  # 保留系统消息和最近8条

                    return result_container["answer"]
                else:
                    logger.error("WebSocket请求失败或无响应")
                    return None
            else:
                logger.error("WebSocket请求超时")
                ws.close()
                return None

        except Exception as e:
            logger.error(f"WebSocket请求异常: {e}")
            return None

    def _get_fallback_response(self, prompt: str) -> str:
        """备用回复函数"""
        user_messages = [msg for msg in self.history if msg["role"] == "user"]
        round_count = len(user_messages)

        logger.info(f"使用备用回复，当前轮次: {round_count}")

        # 根据轮次返回不同问题
        if round_count == 0:
            response = "您好，欢迎参加面试。我看了您的简历，印象深刻。请先简单介绍一下自己，包括您的教育背景和主要工作经验。"
        elif round_count == 1:
            response = "谢谢您的介绍。请详细说说您在工作中遇到的最大挑战是什么，您是如何解决的？"
        elif round_count == 2:
            response = "很好。请谈谈您对这个职位的理解，以及您认为自己最大的优势是什么？"
        else:
            response = "非常感谢您的回答。最后一个问题：您对我们公司或这个职位还有什么想了解的吗？"

        # 更新历史记录
        if prompt:
            self.history.append({"role": "user", "content": prompt})
            self.history.append({"role": "assistant", "content": response})

        return response


# 测试函数
def test_spark_connection():
    """测试星火大模型连接"""
    logger.info("测试星火大模型连接...")

    if not WEBSOCKET_AVAILABLE:
        logger.error("WebSocket不可用，只能使用备用回复")
        return False

    try:
        # 使用简化的系统提示进行测试
        spark = SparkWs(
            SPARK_APPID, SPARK_APIKey, SPARK_APISecret,
            SPARK_URL, SPARK_DOMAIN,
            "你是一个专业的面试官，请用中文回答。"
        )
        response = spark.send_prompt("你好")
        logger.info(f"测试成功，回复长度: {len(response)}")
        logger.info(f"回复内容: {response[:100]}...")
        return True
    except Exception as e:
        logger.error(f"测试失败: {e}")
        return False


if __name__ == "__main__":
    print("🔧 Spark Chat模块测试")
    print("=" * 40)

    # 测试连接
    connection_test = test_spark_connection()
    print(f"连接测试: {'✅' if connection_test else '❌'}")

    print("=" * 40)